%module SerializationModule

%import "LibDL/global_typemaps.i"
%import "LibDL/nn/Modules/Module.i"
%import "LibDL/ordered_dict/ordered_dict.i"

%rename (Load) Serialization::load;
%rename (StateDict) Serialization::state_dict;
%rename (Save) Serialization::save;
%rename (LoadStateDict) Serialization::load_state_dict;
%shared_ptr(nn::Module);

SWIG_JAVABODY_PROXY(public, public, SWIGTYPE)
SWIG_JAVABODY_TYPEWRAPPER(public, public, public, SWIGTYPE)

%{
#include "LibDL/serialize/serialization.h"
#include "./LibDL/serialize/module_keys.h"
%}

%include "./LibDL/serialize/serialization.h"

%template(Tuple) std::pair<std::vector<std::string>, std::vector<std::string>>;
